// Copyright 2006 Bret Taylor
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

#include <iostream>

#include "checker.h"
#include "debug.h"
#include "localization.h"

namespace indie {

// A scoped variable is just a variable with an integer scope that
// indicates its relative position in the AST. We use this to determine
// which variables are available at certain positions in the AST.
class Checker::ScopedVariable {
  public:
    ScopedVariable(const AstVariable& ast, int scope) :
	ast_(ast), scope_(scope) {
    }

    inline const AstVariable& ast() const {
	return ast_;
    }

    inline int scope() const {
	return scope_;
    }

  private:
    const AstVariable& ast_;
    const unsigned int scope_;
};


Checker::Checker(Log& log) :
    log_(log), current_scope_(0) {
    current_class_type_ = NULL;
    current_class_ = NULL;
    current_method_ = NULL;

    // Load built-in classes
    built_ins_[ClassType::Object().name()] = &ClassType::Object();
    built_ins_[ClassType::Int().name()] = &ClassType::Int();
    built_ins_[ClassType::Float().name()] = &ClassType::Float();
    built_ins_[ClassType::Char().name()] = &ClassType::Char();
    built_ins_[ClassType::Boolean().name()] = &ClassType::Boolean();
    built_ins_[ClassType::String().name()] = &ClassType::String();
}

Checker::~Checker() {
}

void Checker::Check(const set<AstFile*>& files) {
    LoadClasses(files);
    CheckCode();
    Cleanup();
}


void Checker::CheckCode() {
    for (map<string, ClassType*>::const_iterator i = classes_.begin();
	 i != classes_.end(); ++i) {
	const ClassType& type(*i->second);
	const AstClass& ast(type.ast());

	// Check all of the methods
	current_class_type_ = &type;
	current_class_ = &ast;
	for (vector<AstMethod*>::const_iterator j = ast.methods().begin();
	     j != ast.methods().end(); ++j) {
	    (*j)->Accept(*this);
	}
	current_class_ = NULL;
	current_class_type_ = NULL;
    }
}


void Checker::LoadClasses(const set<AstFile*>& files) {
    // Load the classes into the map, but don't load superclasses yet
    for (set<AstFile*>::const_iterator i = files.begin(); i != files.end();
	 ++i) {
	const AstFile& file(**i);
	for (vector<AstClass*>::const_iterator j = file.classes().begin();
	     j != file.classes().end(); ++j) {
	    const AstClass& cls(**j);
	    const ClassType* type = LookupClass(cls.identifier());
	    if (type != NULL) {
		if (type->has_ast()) {
		    log_.Error(cls.location(), _("class '%s' already defined "
			       "at %s:%d"), cls.identifier().c_str(),
			       type->ast().location().path().c_str(),
			       type->ast().location().start_line());
		} else {
		    log_.Error(cls.location(), _("class '%s' conflicts with "
			       "built-in type"), cls.identifier().c_str());
		}
		continue;
	    }
	    classes_[cls.identifier()] = new ClassType(cls);
	}
    }

    // Load superclasses, checking for cycles
    for (map<string, ClassType*>::const_iterator i = classes_.begin();
	 i != classes_.end(); ++i) {
	ClassType& type(*i->second);
	const AstClass& ast(type.ast());

	// Load all of the member variables, but make sure they are not
	// in any scope (they have to be referenced with this. or <access>.
	PushScope();
	for (vector<AstVariable*>::const_iterator j = ast.variables().begin();
	     j != ast.variables().end(); ++j) {
	    (*j)->Accept(*this);
	}
	PopScope();

	// Check and initialize the superclass
	if (!ast.has_superclass()) continue;
	const AstClassType& superast(ast.superclass());
	const ClassType* superclass = LookupClass(superast.identifier());
	if (superclass == NULL) {
	    log_.Error(superast.location(), _("unrecognized class "
		       "name '%s'"), superast.identifier().c_str());
	    continue;
	}
	if (superclass->IsSubTypeOf(type)) {
	    log_.Error(superast.location(), _("cyclic class hierarchy "
		       "between '%s' and '%s'"), type.name().c_str(),
		       superclass->name().c_str());
	    continue;
	}
	type.set_superclass(*superclass);
    }

    // Now load all the methods from the "top down", i.e., load superclasses
    // before subclasses. This lets us look for method signature conflicts
    // more easily in subclasses
    set<const ClassType*> loaded;
    for (map<string, ClassType*>::const_iterator i = classes_.begin();
	 i != classes_.end(); ++i) {
	LoadMethods(*i->second, loaded);
    }
}

void Checker::LoadMethods(ClassType& cls, set<const ClassType*>& loaded) {
    // Load our superclas if it is not already loaded
    if (!cls.has_ast()) return;
    if (loaded.find(&cls) != loaded.end()) return;
    if (cls.has_superclass()) {
	ClassType& superclass = const_cast<ClassType&>(cls.superclass());
	LoadMethods(superclass, loaded);
    }

    // Load a vector of our parameter types to construct our method
    // signature
    const AstClass& ast(cls.ast());
    for (vector<AstMethod*>::const_iterator i = ast.methods().begin();
	 i != ast.methods().end(); ++i) {
	const AstMethod& method_ast(**i);
	Method* method;
	if (method_ast.has_return_type()) {
	    const AstUnionType& return_ast(method_ast.return_type());
	    return_ast.Accept(*this);
	    method = new Method(GetTreeType(return_ast));
	} else {
	    method = new Method(UnionType::Empty());
	}

	const vector<AstVariable*>& params(method_ast.parameters());
	for (vector<AstVariable*>::const_iterator j = params.begin();
	     j != params.end(); ++j) {
	    const AstUnionType& param_ast((*j)->type());
	    param_ast.Accept(*this);
	    method->add_parameter(GetTreeType(param_ast));
	}
	cls.add_method(method);
    }
}


void Checker::Cleanup() {
    for (map<string, ClassType*>::iterator i = classes_.begin();
	 i != classes_.end(); ++i) {
	delete i->second;
    }
    classes_.clear();
}


const ClassType* Checker::LookupClass(const string& name) {
    // Search built-ins
    map<string, const ClassType*>::const_iterator bl = built_ins_.find(name);
    if (bl != built_ins_.end()) return bl->second;

    // Search user-defined classes
    map<string, ClassType*>::const_iterator cl = classes_.find(name);
    if (cl == classes_.end()) return NULL;
    return cl->second;
}


void Checker::SetTreeType(const Ast& ast, const UnionType& type) {
    tree_types_[&ast] = &type;
}

const UnionType& Checker::GetTreeType(const Ast& ast) {
    map<const Ast*, const UnionType*>::const_iterator loc =
	tree_types_.find(&ast);
    assert(loc != tree_types_.end());
    return *loc->second;
}


void Checker::SetType(const AstType& ast, const Type& type) {
    types_[&ast] = &type;
}


const Type* Checker::GetType(const AstType& ast) {
    map<const AstType*, const Type*>::const_iterator loc = types_.find(&ast);
    if (loc == types_.end()) return NULL;
    return loc->second;
}


void Checker::PushScope() {
    current_scope_++;
}


void Checker::PopScope() {
    assert(current_scope_ > 0);

    // multimap::iterator is not mutable, so we have to break out of the
    // loop every time we erase a variable. Not very efficient...
    bool more_to_erase = true;
    while (more_to_erase) {
	more_to_erase = false;
	multimap<string, ScopedVariable*>::iterator i;
	for (i = scoped_variables_.begin(); i != scoped_variables_.end();
	     ++i) {
	    ScopedVariable* variable = i->second;
	    assert(variable->scope() <= current_scope_);
	    if (variable->scope() == current_scope_) {
		scoped_variables_.erase(i);
		delete variable;
		more_to_erase = true;
		break;
	    }
	}
    }

    current_scope_--;
}


void Checker::RegisterVariable(const AstVariable& variable) {
    multimap<string, ScopedVariable*>::const_iterator i;
    multimap<string, ScopedVariable*>::const_iterator end =
	scoped_variables_.upper_bound(variable.identifier());
    for (i = scoped_variables_.lower_bound(variable.identifier()); i != end;
	 ++i) {
	const ScopedVariable& other(*i->second);
	if (other.scope() == current_scope_) {
	    log_.Error(variable.location(), _("declaration of '%s' here "
		       "conflicts with declaration on line %d"),
		       variable.identifier().c_str(),
		       other.ast().location().start_line());
	    return;
	}
    }

    ScopedVariable* scoped = new ScopedVariable(variable, current_scope_);
    scoped_variables_.insert(make_pair(variable.identifier(), scoped));
}


const AstVariable* Checker::LookupVariable(const string& name) {
    multimap<string, ScopedVariable*>::const_iterator i =
	scoped_variables_.lower_bound(name);
    multimap<string, ScopedVariable*>::const_iterator end =
	scoped_variables_.upper_bound(name);
    if (i == end) return NULL;

    const ScopedVariable* result = i->second;
    for (++i; i != end; ++i) {
	if (i->second->scope() > result->scope()) {
	    result = i->second;
	}
    }

    return &result->ast();
}


void Checker::Enter(const AstMethod& ast) {
    current_method_ = &ast;
    PushScope();
}


void Checker::Leave(const AstMethod& ast) {
    current_method_ = NULL;
    PopScope();
}


void Checker::Enter(const AstFor& ast) {
    PushScope();
}


void Checker::Leave(const AstFor& ast) {
    PopScope();

    const UnionType& test_type(GetTreeType(ast.test()));
    if (!test_type.Equals(UnionType::Boolean()) &&
	!test_type.Equals(UnionType::Error())) {
	log_.Error(ast.test().location(), _("if test must be boolean "
		   "expression"));
    }
}


void Checker::Enter(const AstBlockStatement& ast) {
    PushScope();
}


void Checker::Leave(const AstBlockStatement& ast) {
    PopScope();
}


void Checker::Leave(const AstVariable& ast) {
    SetTreeType(ast, GetTreeType(ast.type()));
    RegisterVariable(ast);
}


void Checker::Leave(const AstClassType& ast) {
    const Type* type = LookupClass(ast.identifier());
    if (type == NULL) {
	log_.Error(ast.location(), "unrecognized class name '%s'",
		   ast.identifier().c_str());
    } else {
	SetType(ast, *type);
    }
}


void Checker::Leave(const AstArrayType& ast) {
    const Type* element_type = GetType(ast.element_type());
    if (element_type != NULL) {
	// TODO: memory leak
	ArrayType* array_type = new ArrayType(*element_type);
	SetType(ast, *array_type);
    }
}


void Checker::Leave(const AstUnionType& ast) {
    // We may end up walking types multiple times as we load the classes
    // in the first pass
    if (tree_types_.find(&ast) != tree_types_.end()) return;

    // Array and Class types are NULL if they could not be resolved. If
    // any member of our union is unresolved, we set our type to ErrorType.
    auto_ptr<vector<const Type*> > types(new vector<const Type*>);
    for (vector<AstType*>::const_iterator i = ast.types().begin();
	 i != ast.types().end(); ++i) {
	const Type* type = GetType(**i);
	if (type == NULL) {
	    SetTreeType(ast, UnionType::Error());
	    return;
	}
	types->push_back(type);
    }
    // TODO: memory leak
    UnionType* union_type = new UnionType(types.release());
    SetTreeType(ast, *union_type);
}


void Checker::Leave(const AstParenthesizedExpression& ast) {
    SetTreeType(ast, GetTreeType(ast.expression()));
}


void Checker::Leave(const AstAccess& ast) {
    SetTreeType(ast, UnionType::Error());

    if (ast.has_base()) {
	// TODO: Resolve member here
    } else {
	const AstVariable* variable = LookupVariable(ast.identifier());
	if (variable == NULL) {
	    log_.Error(ast.location(), _("undeclared variable '%s'"),
		       ast.identifier().c_str());
	    return;
	}
	SetTreeType(ast, GetTreeType(*variable));
    }
}


void Checker::Leave(const AstArrayAccess& ast) {
    SetTreeType(ast, UnionType::Error());

    auto_ptr<vector<const Type*> > element_types(new vector<const Type*>);
    const UnionType& base_type(GetTreeType(ast.base()));

    // Our type is <error> if the element type is <error>
    if (base_type.Equals(UnionType::Error())) return;

    for (vector<const Type*>::const_iterator i = base_type.types().begin();
	 i != base_type.types().end(); ++i) {
	const ArrayType* array_type = dynamic_cast<const ArrayType*>(*i);
	if (array_type == NULL) {
	    log_.Error(ast.location(), _("invalid array access of non-array "
		       "type '%s'"), (*i)->name().c_str());
	    return;
	}
	element_types->push_back(&array_type->element_type());
    }

    // TODO: memory leak
    UnionType* union_type = new UnionType(element_types.release());
    SetTreeType(ast, *union_type);
}


void Checker::Leave(const AstBinaryExpression& ast) {
    SetTreeType(ast, UnionType::Error());

    const UnionType& left_type(GetTreeType(ast.left_side()));
    const UnionType& right_type(GetTreeType(ast.right_side()));
    vector<const UnionType*> expected_types;
    switch (ast.op()) {
    case Addition:
	expected_types.push_back(&UnionType::Int());
	expected_types.push_back(&UnionType::Float());
	expected_types.push_back(&UnionType::String());
	SetTreeType(ast, left_type);
	break;
    case Subtraction:
    case Multiplication:
    case Division:
    case Mod:
	expected_types.push_back(&UnionType::Int());
	expected_types.push_back(&UnionType::Float());
	SetTreeType(ast, left_type);
	break;
    case Equal:
    case NotEqual:
	expected_types.push_back(&UnionType::Int());
	expected_types.push_back(&UnionType::Float());
	expected_types.push_back(&UnionType::Char());
	expected_types.push_back(&UnionType::Boolean());
	expected_types.push_back(&UnionType::String());
	SetTreeType(ast, UnionType::Boolean());
	break;
    case LessThan:
    case GreaterThan:
    case LessEqual:
    case GreaterEqual:
	expected_types.push_back(&UnionType::Int());
	expected_types.push_back(&UnionType::Float());
	SetTreeType(ast, UnionType::Boolean());
	break;
    case And:
    case Or:
	expected_types.push_back(&UnionType::Boolean());
	SetTreeType(ast, UnionType::Boolean());
	break;
    }

    bool valid_left = left_type.Equals(UnionType::Error());
    bool valid_right = right_type.Equals(UnionType::Error());
    for (vector<const UnionType*>::const_iterator i = expected_types.begin();
	 i != expected_types.end(); ++i) {
	if (!valid_left && left_type.Equals(**i)) valid_left = true;
	if (!valid_right && right_type.Equals(**i)) valid_right = true;
    }
    if (!valid_left) {
	log_.Error(ast.left_side().location(), _("operator not defined on "
		   "type '%s'"), left_type.ToString().c_str());
    }
    if (!valid_right) {
	log_.Error(ast.right_side().location(), _("operator not defined on "
		   "type '%s'"), right_type.ToString().c_str());
    }

    if (!valid_left || !valid_right || left_type.Equals(UnionType::Error()) ||
	right_type.Equals(UnionType::Error())) {
	SetTreeType(ast, UnionType::Error());
	return;
    }

    if (!left_type.Equals(right_type)) {
	log_.Error(ast.location(), _("left side type '%s' does not match "
		   "right side type '%s'"), left_type.ToString().c_str(),
		   right_type.ToString().c_str());
	SetTreeType(ast, UnionType::Error());
	return;
    }
}


void Checker::Leave(const AstUnaryExpression& ast) {
    SetTreeType(ast, UnionType::Error());

    const UnionType& expr_type(GetTreeType(ast.expression()));
    if (expr_type.Equals(UnionType::Error())) return;

    vector<const UnionType*> expected_types;
    switch (ast.op()) {
    case Not:
	expected_types.push_back(&UnionType::Boolean());
	SetTreeType(ast, UnionType::Boolean());
	break;
    case Negative:
	expected_types.push_back(&UnionType::Int());
	expected_types.push_back(&UnionType::Float());
	SetTreeType(ast, expr_type);
	break;
    case PreDecrement:
    case PreIncrement:
    case PostDecrement:
    case PostIncrement:
	expected_types.push_back(&UnionType::Int());
	SetTreeType(ast, UnionType::Int());
	break;
    }

    bool valid = false;
    for (vector<const UnionType*>::const_iterator i = expected_types.begin();
	 i != expected_types.end(); ++i) {
	if (!valid && expr_type.Equals(**i)) valid = true;
    }
    if (!valid) {
	SetTreeType(ast, UnionType::Error());
	log_.Error(ast.location(), _("operator not defined on "
		   "type '%s'"), expr_type.ToString().c_str());
	return;
    }
}


void Checker::Leave(const AstCall& ast) {
    SetTreeType(ast, UnionType::Error());
}


void Checker::Leave(const AstNewInstance& ast) {
    SetTreeType(ast, UnionType::Error());
}


void Checker::Leave(const AstNewArray& ast) {
    SetTreeType(ast, UnionType::Error());
}


void Checker::Leave(const AstTypeTest& ast) {
    SetTreeType(ast, UnionType::Boolean());

    const UnionType& expr_type(GetTreeType(ast.expression()));
    if (expr_type.Equals(UnionType::Error())) return;

    vector<const Type*>* types = new vector<const Type*>;
    types->push_back(GetType(ast.type()));
    UnionType test_type(types);
    if (!test_type.IsSubTypeOf(expr_type)) {
	log_.Error(ast.location(), _("'%s' is not a subtype of '%s'"),
		   test_type.ToString().c_str(), expr_type.ToString().c_str());
    }
}


void Checker::Leave(const AstThis& ast) {
    // TODO: memory leak
    vector<const Type*>* types = new vector<const Type*>;
    types->push_back(&current_class_type());
    UnionType* this_type = new UnionType(types);
    SetTreeType(ast, *this_type);
}


void Checker::Leave(const AstIntegerConstant& ast) {
    SetTreeType(ast, UnionType::Int());
}


void Checker::Leave(const AstStringConstant& ast) {
    SetTreeType(ast, UnionType::String());
}


void Checker::Leave(const AstFloatConstant& ast) {
    SetTreeType(ast, UnionType::Float());
}


void Checker::Leave(const AstBooleanConstant& ast) {
    SetTreeType(ast, UnionType::Boolean());
}


void Checker::Leave(const AstIf& ast) {
    const UnionType& test_type(GetTreeType(ast.test()));
    if (!test_type.Equals(UnionType::Boolean()) &&
	!test_type.Equals(UnionType::Error())) {
	log_.Error(ast.test().location(), _("if test must be boolean "
		   "expression"));
    }
}


void Checker::Leave(const AstWhile& ast) {
    const UnionType& test_type(GetTreeType(ast.test()));
    if (!test_type.Equals(UnionType::Boolean()) &&
	!test_type.Equals(UnionType::Error())) {
	log_.Error(ast.test().location(), _("while test must be boolean "
		   "expression"));
    }
}


void Checker::Leave(const AstAssignment& ast) {
    CheckAssignment(ast.left_side(), ast.right_side());
}


void Checker::Leave(const AstVariableDeclaration& ast) {
    if (ast.has_initial_value()) {
	CheckAssignment(ast.variable(), ast.initial_value());
    }
}


void Checker::CheckAssignment(const Ast& left_side, const Ast& right_side) {
    const UnionType& left_type(GetTreeType(left_side));
    const UnionType& right_type(GetTreeType(right_side));

    if (left_type.Equals(UnionType::Error()) ||
	right_type.Equals(UnionType::Error())) {
	return;
    }

    if (!right_type.IsSubTypeOf(left_type)) {
	log_.Error(right_side.location(), _("cannot assign '%s' to '%s'"),
		   right_type.ToString().c_str(),
		   left_type.ToString().c_str());
    }
}


void Checker::Leave(const AstReturn& ast) {
    if (ast.has_expression()) {
	if (current_method().has_return_type()) {
	    const UnionType& type(GetTreeType(current_method().return_type()));
	    if (!GetTreeType(ast.expression()).IsSubTypeOf(type)) {
		log_.Error(ast.location(), _("return statement must return "
			   "a '%s'"), type.ToString().c_str());
	    }
	} else {
	    log_.Error(ast.location(), _("return value in method with no "
		       "return type"));
	}
    } else if (current_method().has_return_type()) {
	const UnionType& type(GetTreeType(current_method().return_type()));
	log_.Error(ast.location(), _("return statement must return a '%s'"),
		   type.ToString().c_str());
    }
}

}
